import pandas as pd
import numpy as np
import os

def process_dataframes(df_a, df_b):

    # Create a copy of df_a to avoid modifying the original
    result_df = df_a.copy()
    
    # Find rows where tgt is '多个ID'
    mask = result_df['tgt'] == '多个ID'
    
    # Process each row where tgt is '多个ID'
    for idx in result_df[mask].index:
        formula = result_df.loc[idx, 'Formula']
        
        # Find matching rows in df_b
        matches = df_b[df_b['MolecularFormula'] == formula]
        
        # Process each match and create new columns
        for i, (_, match_row) in enumerate(matches.iterrows(), 1):
            # Name column (try NameEn first, then NameZh)
            name_col = f'Name{i}'
            name_value = match_row['NameEN'] if pd.notna(match_row['NameEN']) else match_row['NameZH']
            result_df.loc[idx, name_col] = name_value if pd.notna(name_value) else ''
            
            # Smiles column
            smiles_col = f'Smiles{i}'
            result_df.loc[idx, smiles_col] = match_row['Smiles'] if pd.notna(match_row['Smiles']) else ''
            
            # CAS column (only for the first match)
            CAS_col = f'CAS{i}'
            result_df.loc[idx, CAS_col] = match_row['CAS'] if pd.notna(match_row['CAS']) else ''
    
    return result_df

def load_excel_file(file_path):

    try:
        # Try to read the Excel file
        df = pd.read_excel(file_path, engine='openpyxl')
        return df
    
    except FileNotFoundError:
        raise FileNotFoundError(f"Could not find the file: {file_path}")
    except Exception as e:
        raise Exception(f"Error reading Excel file {file_path}: {str(e)}")

def load_dataframes(path_a, path_b):

    try:
        # Load Excel files
        print(f"Loading first Excel file: {path_a}")
        df_a = load_excel_file(path_a)
        print(f"Loading second Excel file: {path_b}")
        df_b = load_excel_file(path_b)
        
        # Verify required columns exist
        required_cols_a = ['tgt', 'Formula']
        required_cols_b = ['MolecularFormula', 'NameEN', 'NameZH', 'Smiles', 'CAS']
        
        missing_cols_a = [col for col in required_cols_a if col not in df_a.columns]
        missing_cols_b = [col for col in required_cols_b if col not in df_b.columns]
        
        if missing_cols_a:
            raise ValueError(f"Missing required columns in first Excel file: {missing_cols_a}")
        if missing_cols_b:
            raise ValueError(f"Missing required columns in second Excel file: {missing_cols_b}")
        
        return df_a, df_b
    
    except Exception as e:
        print(f"Error loading files: {e}")
        raise

def main():

    path_a = "./processed_file.xlsx"
    path_b = "../../../chem-app/chemapp_resource/pyscript/CreateDB/output_AllTable.xlsx"
    
    try:
        # Load dataframes
        print("Loading Excel files...")
        df_a, df_b = load_dataframes(path_a, path_b)
        
        # Process the dataframes
        print("Processing dataframes...")
        result = process_dataframes(df_a, df_b)
        
        # Save the result as Excel file
        output_path = "Matched_DB.xlsx"
        result.to_excel(output_path, index=False, engine='openpyxl')
        print(f"Processing completed successfully! Results saved to {output_path}")
        
    except Exception as e:
        print(f"An error occurred: {e}")

if __name__ == "__main__":
    main()